#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
UCEIS模型训练脚本 - UC Model专用版本
专门在uc_model目录内运行，结果保存在uc_model内部
"""

import os
import sys
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import numpy as np
from sklearn.metrics import classification_report, balanced_accuracy_score, accuracy_score
from sklearn.metrics import confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
from pathlib import Path
import json
import time
from datetime import datetime
import argparse
from tqdm import tqdm

# 获取uc_model目录路径
UC_MODEL_DIR = Path(__file__).parent
sys.path.insert(0, str(UC_MODEL_DIR))

# 直接导入，不需要uc_model前缀
from utils.dataset import UCEISDataset, train_transform, val_transform, custom_collate_fn, split_dataset
from models.fusion_model import EndoMultimodalModel


class UCEISTrainer:
    """UCEIS模型训练器"""

    def __init__(self, config):
        self.config = config
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # 设置随机种子
        torch.manual_seed(config.seed)
        np.random.seed(config.seed)

        # 初始化模型、数据、优化器等
        self.setup_model()
        self.setup_data()
        self.setup_optimizer()
        self.setup_logging()

        # 评估指标历史
        self.train_history = {
            'loss': [],
            'accuracy': [],
            'balanced_accuracy': [],
            'f1_macro': []
        }
        self.val_history = {
            'loss': [],
            'accuracy': [],
            'balanced_accuracy': [],
            'f1_macro': []
        }

        # 最佳模型指标
        self.best_val_f1 = 0.0
        self.best_epoch = 0

    def setup_model(self):
        """初始化模型"""
        print(f"正在初始化模型，使用设备: {self.device}")

        self.model = EndoMultimodalModel(
            image_pretrained=self.config.image_pretrained,
            text_pretrained=self.config.text_pretrained,
            feature_dim=self.config.feature_dim,
            num_classes=self.config.num_classes
        ).to(self.device)

        print(f"模型参数数量: {sum(p.numel() for p in self.model.parameters()):,}")
        print(f"可训练参数数量: {sum(p.numel() for p in self.model.parameters() if p.requires_grad):,}")

        # 损失函数
        self.criterion = nn.CrossEntropyLoss(weight=self.get_class_weights())

    def get_class_weights(self):
        """计算类别权重以处理不平衡数据"""
        # 这里应该基于实际数据分布计算，暂时使用平衡权重
        class_weights = torch.ones(self.config.num_classes)
        return class_weights.to(self.device)

    def setup_data(self):
        """初始化数据集和数据加载器"""
        print(f"正在加载数据集: {self.config.data_dir}")

        # 创建完整数据集
        full_dataset = UCEISDataset(
            data_dir=self.config.data_dir,
            transform=train_transform  # 先使用训练变换
        )

        # 分割数据集
        train_dataset, val_dataset = split_dataset(
            full_dataset,
            train_ratio=self.config.train_ratio,
            random_seed=self.config.seed
        )

        # 为验证集设置验证变换
        val_dataset.dataset.transform = val_transform

        print(f"训练集样本数: {len(train_dataset)}")
        print(f"验证集样本数: {len(val_dataset)}")

        # 创建数据加载器
        self.train_loader = DataLoader(
            train_dataset,
            batch_size=self.config.batch_size,
            shuffle=True,
            num_workers=self.config.num_workers,
            pin_memory=False,  # 避免CUDA tensor pin_memory问题
            collate_fn=lambda batch: custom_collate_fn(batch, self.device)
        )

        self.val_loader = DataLoader(
            val_dataset,
            batch_size=self.config.batch_size,
            shuffle=False,
            num_workers=self.config.num_workers,
            pin_memory=False,  # 避免CUDA tensor pin_memory问题
            collate_fn=lambda batch: custom_collate_fn(batch, self.device)
        )

    def setup_optimizer(self):
        """初始化优化器和学习率调度器"""
        # 分层学习率设置
        param_groups = [
            {
                'params': self.model.image_branch.parameters(),
                'lr': self.config.lr * 0.1  # 图像分支使用较小学习率
            },
            {
                'params': self.model.text_branch.parameters(),
                'lr': self.config.lr * 0.1  # 文本分支使用较小学习率
            },
            {
                'params': self.model.fusion_layers.parameters(),
                'lr': self.config.lr
            },
            {
                'params': self.model.classifier.parameters(),
                'lr': self.config.lr
            }
        ]

        self.optimizer = optim.AdamW(
            param_groups,
            weight_decay=self.config.weight_decay
        )

        # 学习率调度器
        self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer,
            mode='max',
            factor=0.5,
            patience=self.config.lr_patience
        )

    def setup_logging(self):
        """设置日志记录"""
        # 创建输出目录
        self.output_dir = Path(self.config.output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # 保存配置
        config_path = self.output_dir / 'config.json'
        with open(config_path, 'w', encoding='utf-8') as f:
            json.dump(vars(self.config), f, indent=2, ensure_ascii=False)

        print(f"训练输出目录: {self.output_dir}")

    def train_one_epoch(self, epoch):
        """训练一个epoch"""
        self.model.train()
        epoch_loss = 0.0
        all_predictions = []
        all_labels = []

        progress_bar = tqdm(self.train_loader, desc=f'Epoch {epoch+1}/{self.config.epochs} [Train]')

        for batch_idx, batch in enumerate(progress_bar):
            images = batch['images']
            text_inputs = batch['text_inputs']
            labels = batch['labels']

            # 清零梯度
            self.optimizer.zero_grad()

            # 前向传播
            outputs = self.model(images, text_inputs)
            loss = self.criterion(outputs, labels)

            # 反向传播
            loss.backward()

            # 梯度裁剪
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)

            # 参数更新
            self.optimizer.step()

            # 统计
            epoch_loss += loss.item()
            predictions = torch.argmax(outputs, dim=1)
            all_predictions.extend(predictions.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

            # 更新进度条
            current_lr = self.optimizer.param_groups[0]['lr']
            progress_bar.set_postfix({
                'loss': f'{loss.item():.4f}',
                'lr': f'{current_lr:.6f}'
            })

        # 计算平均损失和准确率
        avg_loss = epoch_loss / len(self.train_loader)
        accuracy = accuracy_score(all_labels, all_predictions)
        balanced_acc = balanced_accuracy_score(all_labels, all_predictions)

        # 8分类F1分数
        f1_macro = self.calculate_f1_macro(all_labels, all_predictions)

        # 记录历史
        self.train_history['loss'].append(avg_loss)
        self.train_history['accuracy'].append(accuracy)
        self.train_history['balanced_accuracy'].append(balanced_acc)
        self.train_history['f1_macro'].append(f1_macro)

        print(f'Train - Loss: {avg_loss:.4f}, Acc: {accuracy:.4f}, Balanced Acc: {balanced_acc:.4f}, F1-macro: {f1_macro:.4f}')

        return avg_loss, accuracy, balanced_acc, f1_macro

    def evaluate(self, epoch):
        """验证模型"""
        self.model.eval()
        epoch_loss = 0.0
        all_predictions = []
        all_labels = []

        with torch.no_grad():
            progress_bar = tqdm(self.val_loader, desc=f'Epoch {epoch+1}/{self.config.epochs} [Val]')

            for batch in progress_bar:
                images = batch['images']
                text_inputs = batch['text_inputs']
                labels = batch['labels']

                # 前向传播
                outputs = self.model(images, text_inputs)
                loss = self.criterion(outputs, labels)

                # 统计
                epoch_loss += loss.item()
                predictions = torch.argmax(outputs, dim=1)
                all_predictions.extend(predictions.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())

                progress_bar.set_postfix({'loss': f'{loss.item():.4f}'})

        # 计算指标
        avg_loss = epoch_loss / len(self.val_loader)
        accuracy = accuracy_score(all_labels, all_predictions)
        balanced_acc = balanced_accuracy_score(all_labels, all_predictions)

        # 8分类评估
        f1_macro = self.calculate_f1_macro(all_labels, all_predictions)

        # 3分类映射评估
        f1_macro_3class, detailed_metrics_3class = self.evaluate_3class_mapping(
            all_labels, all_predictions
        )

        # 记录历史
        self.val_history['loss'].append(avg_loss)
        self.val_history['accuracy'].append(accuracy)
        self.val_history['balanced_accuracy'].append(balanced_acc)
        self.val_history['f1_macro'].append(f1_macro)

        print(f'Val - Loss: {avg_loss:.4f}, Acc: {accuracy:.4f}, Balanced Acc: {balanced_acc:.4f}')
        print(f'8-class F1-macro: {f1_macro:.4f}, 3-class F1-macro: {f1_macro_3class:.4f}')

        # 打印详细3分类指标
        print("\n3分类详细指标:")
        for class_name, metrics in detailed_metrics_3class.items():
            print(f"  {class_name}: P={metrics['precision']:.3f}, R={metrics['recall']:.3f}, F1={metrics['f1-score']:.3f}")

        return avg_loss, accuracy, balanced_acc, f1_macro, f1_macro_3class

    def map_8_to_3_class(self, score_8_class_idx):
        """
        将8分类索引映射到3分类
        输入: 0-7 (对应1-8分)
        输出: 0-2 (轻度/中度/重度)
        """
        if score_8_class_idx <= 2:  # 1-3分 -> 轻度
            return 0
        elif score_8_class_idx <= 5:  # 4-6分 -> 中度
            return 1
        else:  # 7-8分 -> 重度
            return 2

    def evaluate_3class_mapping(self, labels_true, labels_pred):
        """评估3分类映射结果"""
        # 转换为3分类
        labels_true_3class = [self.map_8_to_3_class(label) for label in labels_true]
        labels_pred_3class = [self.map_8_to_3_class(label) for label in labels_pred]

        # 类别名称
        class_names = ['轻度(1-3分)', '中度(4-6分)', '重度(7-8分)']

        # 确保包含所有可能的类别
        unique_labels = sorted(set(labels_true_3class + labels_pred_3class))
        labels = list(range(3))  # 0, 1, 2 代表三种严重程度

        # 生成分类报告
        report = classification_report(
            labels_true_3class,
            labels_pred_3class,
            labels=labels,
            target_names=class_names,
            output_dict=True,
            zero_division=0
        )

        f1_macro = report['macro avg']['f1-score']

        detailed_metrics = {}
        for i, class_name in enumerate(class_names):
            if str(i) in report:
                detailed_metrics[class_name] = report[str(i)]
            else:
                # 如果某个类别不存在，添加默认值
                detailed_metrics[class_name] = {
                    'precision': 0.0,
                    'recall': 0.0,
                    'f1-score': 0.0,
                    'support': 0
                }

        return f1_macro, detailed_metrics

    def calculate_f1_macro(self, labels_true, labels_pred):
        """计算F1-macro分数"""
        report = classification_report(
            labels_true, labels_pred,
            output_dict=True,
            zero_division=0
        )
        return report['macro avg']['f1-score']

    def save_checkpoint(self, epoch, is_best=False):
        """保存模型检查点"""
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict(),
            'train_history': self.train_history,
            'val_history': self.val_history,
            'best_val_f1': self.best_val_f1,
            'config': vars(self.config)
        }

        # 保存最新检查点
        checkpoint_path = self.output_dir / 'latest_model.pth'
        torch.save(checkpoint, checkpoint_path)

        # 保存最佳模型到weights目录
        if is_best:
            best_path = self.output_dir / 'best_model.pth'
            torch.save(checkpoint, best_path)
            print(f"保存最佳模型: {best_path}")

            # 同时保存到weights目录
            weights_dir = UC_MODEL_DIR / 'weights'
            weights_best_path = weights_dir / 'best_model.pth'
            torch.save(checkpoint, weights_best_path)
            print(f"复制最佳模型到weights目录: {weights_best_path}")

    def save_training_summary(self):
        """保存训练总结"""
        summary = {
            'training_completed': datetime.now().isoformat(),
            'total_epochs': self.config.epochs,
            'best_epoch': self.best_epoch,
            'best_val_f1': self.best_val_f1,
            'final_train_loss': self.train_history['loss'][-1] if self.train_history['loss'] else None,
            'final_val_loss': self.val_history['loss'][-1] if self.val_history['loss'] else None,
            'final_train_accuracy': self.train_history['accuracy'][-1] if self.train_history['accuracy'] else None,
            'final_val_accuracy': self.val_history['accuracy'][-1] if self.val_history['accuracy'] else None,
            'device': str(self.device),
            'model_parameters': sum(p.numel() for p in self.model.parameters()),
            'trainable_parameters': sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        }

        # 保存到results目录
        results_path = UC_MODEL_DIR / 'results' / 'training_summary.json'
        with open(results_path, 'w', encoding='utf-8') as f:
            json.dump(summary, f, indent=2, ensure_ascii=False)

        print(f"训练总结已保存: {results_path}")

    def train(self):
        """主训练循环"""
        print(f"开始训练，共 {self.config.epochs} 个epoch")
        print(f"训练模式: {'测试版' if self.config.test_mode else '标准版'}")

        start_time = time.time()

        for epoch in range(self.config.epochs):
            print(f"\n{'='*50}")
            print(f"Epoch {epoch+1}/{self.config.epochs}")
            print(f"{'='*50}")

            # 训练阶段
            train_loss, train_acc, train_balanced_acc, train_f1 = self.train_one_epoch(epoch)

            # 验证阶段
            val_loss, val_acc, val_balanced_acc, val_f1, val_f1_3class = self.evaluate(epoch)

            # 学习率调度
            self.scheduler.step(val_f1_3class)  # 基于3分类F1调度

            # 保存检查点
            is_best = val_f1_3class > self.best_val_f1
            if is_best:
                self.best_val_f1 = val_f1_3class
                self.best_epoch = epoch

            self.save_checkpoint(epoch, is_best)

            # 打印当前最佳结果
            print(f"\n当前最佳: Epoch {self.best_epoch+1}, 3-class F1: {self.best_val_f1:.4f}")

        training_time = time.time() - start_time
        print(f"\n{'='*50}")
        print(f"训练完成！总耗时: {training_time/60:.1f} 分钟")
        print(f"最佳模型: Epoch {self.best_epoch+1}, 3-class F1: {self.best_val_f1:.4f}")
        print(f"{'='*50}")

        # 保存训练总结
        self.save_training_summary()

        return self.best_val_f1


class Config:
    """训练配置"""
    def __init__(self, test_mode=False):
        # 基本配置
        self.test_mode = test_mode
        self.seed = 42
        self.data_dir = r"D:\肠内镜数据库\UCEIS1-8"

        # 修改输出目录到uc_model内部
        uc_model_dir = Path(__file__).parent
        if test_mode:
            self.output_dir = str(uc_model_dir / "outputs" / "test_mode")
        else:
            self.output_dir = str(uc_model_dir / "outputs" / "standard_training")

        # 模型配置
        self.image_pretrained = True
        self.text_pretrained = 'bert-base-uncased'
        self.feature_dim = 256
        self.num_classes = 8

        # 训练配置
        if test_mode:
            # 测试版配置（30分钟内完成）
            self.epochs = 5
            self.batch_size = 16
            self.lr = 5e-4
            self.weight_decay = 1e-3
            self.lr_patience = 2
        else:
            # 标准版配置（完整训练）
            self.epochs = 50
            self.batch_size = 32
            self.lr = 2e-4
            self.weight_decay = 1e-4
            self.lr_patience = 5

        # 数据配置
        self.train_ratio = 0.8
        self.num_workers = 0  # 避免multiprocessing问题


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='UCEIS多模态模型训练 - UC Model专用版本')
    parser.add_argument('--test', action='store_true', help='测试版训练（30分钟）')
    parser.add_argument('--standard', action='store_true', help='标准版训练（完整）')
    args = parser.parse_args()

    # 确定训练模式
    if args.test:
        print("启动测试版训练...")
        config = Config(test_mode=True)
    elif args.standard:
        print("启动标准版训练...")
        config = Config(test_mode=False)
    else:
        print("请选择训练模式: --test (测试版) 或 --standard (标准版)")
        return

    # 检查CUDA
    if torch.cuda.is_available():
        print(f"使用GPU: {torch.cuda.get_device_name()}")
        print(f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
    else:
        print("使用CPU训练（速度较慢）")

    # 创建训练器并开始训练
    trainer = UCEISTrainer(config)
    best_f1 = trainer.train()

    print(f"\n训练完成！最佳3分类F1分数: {best_f1:.4f}")


if __name__ == "__main__":
    main()